from io import BytesIO
import random
import string
from typing import Tuple
from fastapi import HTTPException
from PIL import Image, ImageDraw, ImageFont
import base64

# 验证码配置
class CaptchaConfig:
    def __init__(self):
        self.width = 120
        self.height = 40
        self.font_size = 25
        self.code_length = 4
        self.line_num = 3
        self.background_color = (255, 255, 255)
        self.font_path = None  # 可以使用默认字体或指定字体文件路径

config = CaptchaConfig()

def generate_random_code(length: int = 4) -> str:
    """生成随机验证码"""
    chars = string.ascii_letters + string.digits
    return ''.join(random.choice(chars) for _ in range(length))

def generate_captcha_image(code: str) -> bytes:
    """生成验证码图片"""
    # 创建图片
    image = Image.new('RGB', (config.width, config.height), config.background_color)
    draw = ImageDraw.Draw(image)
    
    # 绘制干扰线
    for _ in range(config.line_num):
        begin = (random.randint(0, config.width), random.randint(0, config.height))
        end = (random.randint(0, config.width), random.randint(0, config.height))
        draw.line([begin, end], fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    
    # 绘制验证码
    try:
        font = ImageFont.truetype(config.font_path, config.font_size)
    except:
        font = ImageFont.load_default()
    
    for i, char in enumerate(code):
        draw.text(
            (10 + i * 30, 5), 
            char, 
            font=font, 
            fill=(random.randint(0, 160), random.randint(0, 160), random.randint(0, 160))
        )
    
    # 输出图片
    buffer = BytesIO()
    image.save(buffer, format='JPEG')
    return buffer.getvalue()

def generate_captcha() -> Tuple[str, str]:
    """生成验证码和图片base64编码"""
    code = generate_random_code(config.code_length)
    image_data = generate_captcha_image(code)
    image_base64 = base64.b64encode(image_data).decode('utf-8')
    return code, f"data:image/jpeg;base64,{image_base64}"